package com.deepseek.project.websocket;

import com.deepseek.project.constant.Constants;
import com.deepseek.project.model.TalkInfoDeepSeekHistory;
import com.deepseek.project.service.ITalkInfoService;
import com.deepseek.project.tool.AbstractChatTool;
import com.deepseek.project.tool.DeepSeekLocalTool;
import com.deepseek.project.tool.DeepSeekOnlineTool;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.lang.reflect.Type;
import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * author: hulei
 * WebSocket处理器，用于处理与客户端的WebSocket通信
 */
@Slf4j
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    // 使用talkInfoId作为键的会话映射
    private final Map<String, AbstractChatTool> talkSessionMap = new ConcurrentHashMap<>();
    private final ExecutorService executorService = Executors.newFixedThreadPool(10);

    private static final Gson gson = new Gson();
    private static final ObjectMapper objectMapper = new ObjectMapper();

    @Resource
    private ITalkInfoService italkInfoService;

    @Value("${deepseek.tool.type:local}")
    private String deepSeekToolType;

    @Override
    public void afterConnectionEstablished(@NotNull WebSocketSession session) {
        String talkInfoId = extractTalkInfoId(session);
        if (talkInfoId == null) {
            closeWithError(session);
            return;
        }
        // 创建或获取已有会话工具
        talkSessionMap.computeIfAbsent(talkInfoId, id -> {
            log.info("构建新会话工具,talkInfoId: {}", talkInfoId);
            List<Map<String, String>> conversationHistory = new ArrayList<>();
            if (!talkInfoId.startsWith(Constants.TEMP_TALK_INFO_ID_PREFIX)) {
                //根据会话框id查询会话框历史记录
                List<TalkInfoDeepSeekHistory> historyList = italkInfoService.getDeepSeekHistoryListByTalkInfoId(Integer.parseInt(talkInfoId));
                for (TalkInfoDeepSeekHistory history : historyList) {
                    Type type = new TypeToken<Map<String, String>>() {
                    }.getType();
                    Map<String, String> messageMap = gson.fromJson(history.getContent(), type);
                    conversationHistory.add(messageMap);
                }
            }
            return createDeepSeekTool(conversationHistory);
        });
    }

    /**
     * 创建AbstractDeepSeekTool具体类的实例，根据配置的toolType参数来选择创建哪种类型的工具
     */
    private AbstractChatTool createDeepSeekTool(List<Map<String, String>> conversationHistory) {
        if ("online".equalsIgnoreCase(deepSeekToolType)) {
            return new DeepSeekOnlineTool(italkInfoService, conversationHistory);
        } else {
            return new DeepSeekLocalTool(italkInfoService, conversationHistory);
        }
    }

    /**
     * 从WebSocketSession中提取talkInfoId参数值
     * 此方法主要用于从WebSocket连接的URI中提取出talkInfoId参数值，该参数值用于标识对话信息
     * 如果URI为空，或者没有找到talkInfoId参数，则返回null
     *
     * @param session WebSocketSession对象，包含客户端与服务器之间的WebSocket连接信息
     * @return String 返回提取出的talkInfoId参数值，如果没有找到则返回null
     */
    private String extractTalkInfoId(WebSocketSession session) {
        try {
            URI uri = session.getUri();
            if (uri == null) return null;

            return Arrays.stream(uri.getQuery().split("&"))
                    .filter(param -> param.startsWith("talkInfoId="))
                    .findFirst()
                    .map(param -> param.split("=")[1])
                    .orElse(null);
        } catch (Exception e) {
            log.error("解析talkInfoId失败", e);
            return null;
        }
    }

    /**
     * 处理WebSocket消息
     */
    @Override
    protected void handleTextMessage(@NotNull WebSocketSession session, @NotNull TextMessage message) {
        try {
            log.info("收到消息：{}", message.getPayload());
            JsonNode json = objectMapper.readTree(message.getPayload());
            String talkInfoId = json.get("talkInfoId").asText();
            String content = json.get("content").asText();
            AbstractChatTool tool = talkSessionMap.get(talkInfoId);
            if (tool == null) {
                log.warn("找不到对应的会话工具，talkInfoId: {}", talkInfoId);
                sendErrorMessage(session, "无效 session");
                return;
            }
            log.info("成功获取到会话工具，talkInfoId: {}", talkInfoId);
            if ("stopSending".equalsIgnoreCase(content)) {
                handleStopCommand(tool, session);
            } else {
                handleNormalMessage(tool, session, content, talkInfoId);
            }
        } catch (IOException e) {
            log.error("消息解析失败", e);
            sendErrorMessage(session, "Invalid message format");
        }
    }

    /**
     * 处理前端点击发送的停止生成命令
     */
    private void handleStopCommand(AbstractChatTool tool, WebSocketSession session) {
        tool.setStopSending(true);
        try {
            session.sendMessage(new TextMessage("[END_OF_MESSAGE_GENERATE]"));
        } catch (IOException e) {
            log.error("停止命令响应失败", e);
        }
    }

    /**
     * 处理普通消息
     */
    private void handleNormalMessage(AbstractChatTool tool, WebSocketSession session, String content, String talkInfoId) {
        tool.setStopSending(false);
        executorService.submit(() -> {
            try {
                tool.processMessage(content, session, talkInfoId);
            } catch (Exception e) {
                log.error("消息处理异常", e);
                sendErrorMessage(session, "消息处理发生异常");
            }
        });
    }

    @Override
    public void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus status) {
        String talkInfoId = extractTalkInfoId(session);
        if (talkInfoId != null) {
            // 根据业务需求决定是否立即清理资源
            // 如果是持久化会话可以保留，临时会话则移除,因为每次刷新页面临时会话都会新建，原来的就没用了，需要移除
            if (talkInfoId.startsWith(Constants.TEMP_TALK_INFO_ID_PREFIX)) {
                talkSessionMap.remove(talkInfoId);
            }
            log.info("会话关闭，talkInfoId: {}", talkInfoId);
        }
    }

    /**
     * 发送错误消息
     */
    private void sendErrorMessage(WebSocketSession session, String error) {
        try {
            session.sendMessage(new TextMessage(
                    objectMapper.createObjectNode()
                            .put("type", "error")
                            .put("message", error)
                            .toString()
            ));
        } catch (IOException e) {
            log.error("错误信息发送失败", e);
        }
    }

    private void closeWithError(WebSocketSession session) {
        try {
            session.close(new CloseStatus(CloseStatus.BAD_DATA.getCode(), "缺失talkInfoId会话框ID参数"));
        } catch (IOException e) {
            log.error("关闭连接失败", e);
        }
    }
}
